#include <array>
#include <iostream>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define MSG_LEN 16
#define MSG_CNT 10

#define AES_DIGEST_LEN MSG_LEN
#define INPUT_SIZE (MSG_LEN * MSG_CNT)
#define OUTPUT_SIZE (AES_DIGEST_LEN * MSG_CNT)
#define KEY_SIZE 32

enum class AES_TYPE { AES128 = 16, AES192 = 24, AES256 = 32 };

#ifndef WAFER_OPENSSL
extern "C" {
void CUDA(MemRef<unsigned char, 1> *inputs, int msgLen,
          MemRef<unsigned char, 1> *outputs, int cptLen,
          MemRef<unsigned char, 1> *key, int keyLen, int msgCnt);
}

void AES_CRACK_OP(MemRef<unsigned char, 1> *inputs, int msgLen,
                  MemRef<unsigned char, 1> *outputs, int cptLen,
                  MemRef<unsigned char, 1> *key, int keyLen, int msgCnt) {
  CUDA(inputs, msgLen, outputs, cptLen, key, keyLen, msgCnt);
}
#else
#define OPENSSL_CPU_NUM 2
extern "C" uint64_t sstGetCpuId();
#endif

uint8_t exampleInput[INPUT_SIZE] = {
    0xE2, 0x4C, 0xF1, 0xEC, 0x3C, 0x18, 0xAD, 0x6E, 0x8E, 0xB8, 0x45, 0xEA,
    0xD1, 0x34, 0xAF, 0xE2, 0xD5, 0xAB, 0xCA, 0x17, 0xC2, 0x45, 0x04, 0x02,
    0x76, 0x28, 0xC5, 0x57, 0x05, 0xC6, 0x46, 0x38, 0x45, 0xD2, 0x00, 0xE3,
    0xB1, 0xD9, 0x28, 0x6B, 0x92, 0x7D, 0xEC, 0xA0, 0xB6, 0x56, 0xFA, 0x92,
    0x88, 0xB5, 0x96, 0x2B, 0xB3, 0x60, 0xF4, 0x2A, 0x83, 0x67, 0x6C, 0x21,
    0x5B, 0xC8, 0x61, 0x5B, 0x96, 0x7E, 0x12, 0x39, 0x06, 0xC8, 0xB5, 0xD3,
    0x0F, 0xC2, 0x25, 0xB9, 0xD1, 0x39, 0x67, 0x5A, 0x69, 0x9D, 0xD6, 0x07,
    0x0B, 0xBF, 0x66, 0xD6, 0xAF, 0xD1, 0x94, 0x42, 0x2F, 0x47, 0x9F, 0x04,
    0xCA, 0x85, 0xC7, 0x28, 0x36, 0xC1, 0x86, 0x48, 0xF4, 0x05, 0xD9, 0xBE,
    0x44, 0x3F, 0x30, 0x8C, 0xFA, 0x2E, 0xB6, 0x4E, 0x65, 0x56, 0x15, 0x58,
    0x56, 0x62, 0x19, 0x5D, 0xD7, 0xFD, 0xE9, 0x66, 0x59, 0x0E, 0x4B, 0xDD,
    0xE9, 0x39, 0x3D, 0x10, 0x55, 0xE1, 0xA5, 0xB9, 0x0B, 0x42, 0x0B, 0x39,
    0x8D, 0x52, 0x63, 0x26, 0x52, 0xBF, 0x8D, 0xA0, 0xF0, 0xC1, 0x46, 0xB0,
    0x9A, 0x91, 0x17, 0x7E};
uint8_t exampleKey[KEY_SIZE] = {
    'A', '0', 'B', '1', 'C', '2', 'D', '3', 'E', '4', 'F',
    '5', 'G', '6', 'H', '7', 'I', '8', 'J', '9', 'K', '0',
    'L', '1', 'M', '2', 'N', '3', 'O', '4', 'P', '5',
};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> key{exampleKey, KEY_SIZE};
uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

void printResult(uint8_t *output) {
  // Print the output.
  std::cout << "[";
  for (size_t i = 0; i < AES_DIGEST_LEN; i++) {
    printf("%c", output[i]);
  }
  std::cout << "]" << std::endl;
}

void printInput(int idx, uint8_t *input) {
  // Print the input.
  std::cout << "Input " << idx << ": [0x";
  for (size_t i = 0; i < MSG_LEN; i++) {
    printf("%02X", input[i]);
  }
  std::cout << "] ==> ";
}

void generateResultGPUAES() {
  int keyLen = static_cast<int>(AES_TYPE::AES128);

#ifdef WAFER_OPENSSL
  uint64_t cpuid = sstGetCpuId();
  uint64_t step = MSG_NUM / OPENSSL_CPU_NUM;
  uint64_t start = cpuid * step;
  uint64_t end =
      (cpuid == OPENSSL_CPU_NUM - 1) ? MSG_NUM : ((cpuid + 1) * step);

  const clock_t begin_time = clock();
  for (int i = start; i < end; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));

    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

    const EVP_CIPHER *cipher;
    if (keyLen == 16)
      cipher = EVP_aes_128_ecb();
    else if (keyLen == 24)
      cipher = EVP_aes_192_ecb();
    else // keyLen == 32
      cipher = EVP_aes_256_ecb();

    EVP_DecryptInit_ex(ctx, cipher, NULL, exampleKey, NULL);
    int outlen;
    int plaintext_len = 0;
    EVP_DecryptUpdate(ctx, exampleOutput + (i * AES_DIGEST_LEN), &outlen,
                      exampleInput + (i * MSG_LEN), AES_DIGEST_LEN);
    plaintext_len += outlen;
    EVP_DecryptFinal_ex(
        ctx, exampleOutput + (i * AES_DIGEST_LEN) + plaintext_len, &outlen);
    plaintext_len += outlen;

    printResult(exampleOutput + (i * AES_DIGEST_LEN));
    EVP_CIPHER_CTX_free(ctx);
  }

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "cpuid: " << cpuid << std::endl << "handle data: [" << start << ", " << end << "]" << std::endl <<"Calculate duration: " << std::fixed << duration << std::endl;
#else
  const clock_t begin_time = clock();
#pragma wafer parallel
  AES_CRACK_OP(&input, MSG_LEN, &output, AES_DIGEST_LEN, &key, keyLen, MSG_CNT);
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;

  std::cout << "\n--------------------------------------------------------"
               "\nMLIR AES\n";

  for (int i = 0; i < MSG_CNT; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));
    printResult(exampleOutput + (i * AES_DIGEST_LEN));
  }

  std::cout << "Key " << ": [";
  for (size_t i = 0; i < keyLen; i++) {
    std::cout << key[i];
  }
  std::cout << "]" << std::endl;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

int main() {
  generateResultGPUAES();
  return 0;
}
